{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Understanding_TensorFlow_Distributions_Shapes.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [
        "wDRB80oLrG48"
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FVKYfpQVYPaJ",
        "colab_type": "text"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "htHLjlnLYSoB",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DcriL2xPrG3_",
        "colab_type": "text"
      },
      "source": [
        "# Understanding TensorFlow Distributions Shapes\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/probability/blob/master/tensorflow_probability/examples/jupyter_notebooks/Understanding_TensorFlow_Distributions_Shapes.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/probability/blob/master/tensorflow_probability/examples/jupyter_notebooks/Understanding_TensorFlow_Distributions_Shapes.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J6t0EUihrG4B",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import collections\n",
        "\n",
        "import tensorflow as tf\n",
        "import tensorflow_probability as tfp\n",
        "tfd = tfp.distributions\n",
        "tfb = tfp.bijectors\n",
        "\n",
        "tf.compat.v1.enable_eager_execution()\n",
        "\n",
        "from __future__ import print_function"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QD5lzFZerG4H",
        "colab_type": "text"
      },
      "source": [
        "## Basics\n",
        "\n",
        "There are three important concepts associated with TensorFlow Distributions shapes:\n",
        "- *Event shape* describes the shape of a single draw from the distribution; it may be dependent across dimensions. For scalar distributions, the event shape is `[]`. For a 5-dimensional MultivariateNormal, the event shape is `[5]`.\n",
        "- *Batch shape* describes independent, not identically distributed draws, aka a \"batch\" of distributions.\n",
        "- *Sample shape* describes independent, identically distributed draws of batches from the distribution family.\n",
        "\n",
        "The event shape and the batch shape are properties of a `Distribution` object, whereas the sample shape is associated with a specific call to `sample` or `log_prob`.\n",
        "\n",
        "This notebook's purpose is to illustrate these concepts through examples, so if this isn't immediately obvious, don't worry!\n",
        "\n",
        "For another conceptual overview of these concepts, see [this blog post](https://ericmjl.github.io/blog/2019/5/29/reasoning-about-shapes-and-probability-distributions/)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yU34kIHDrG4I",
        "colab_type": "text"
      },
      "source": [
        "### A note on TensorFlow Eager.\n",
        "\n",
        "This entire notebook is written using [TensorFlow Eager](https://research.googleblog.com/2017/10/eager-execution-imperative-define-by.html). None of the concepts presented *rely* on Eager, although with Eager, distribution batch and event shapes are evaluated (and therefore known) when the `Distribution` object is created in Python, whereas in graph (non-Eager mode), it is possible to define distributions whose event and batch shapes are undetermined until the graph is run."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MeirD-0JrG4K",
        "colab_type": "text"
      },
      "source": [
        "## Scalar Distributions\n",
        "\n",
        "As we noted above, a `Distribution` object has defined event and batch shapes. We'll start with a utility to describe distributions:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bq8guNPtrG4M",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def describe_distributions(distributions):\n",
        "  print('\\n'.join([str(d) for d in distributions]))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "06CafVXWrG4Q",
        "colab_type": "text"
      },
      "source": [
        "In this section we'll explore *scalar* distributions: distributions with an event shape of `[]`. A typical example is the Poisson distribution, specified by a `rate`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Sdz1OMg7rG4S",
        "colab_type": "code",
        "outputId": "117e7ab4-7f01-4380-e0ba-faae53ec0c3c",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 102
        }
      },
      "source": [
        "poisson_distributions = [\n",
        "    tfd.Poisson(rate=1., name='One Poisson Scalar Batch'),\n",
        "    tfd.Poisson(rate=[1., 10., 100.], name='Three Poissons'),\n",
        "    tfd.Poisson(rate=[[1., 10., 100.,], [2., 20., 200.]],\n",
        "                name='Two-by-Three Poissons'),\n",
        "    tfd.Poisson(rate=[1.], name='One Poisson Vector Batch'),\n",
        "    tfd.Poisson(rate=[[1.]], name='One Poisson Expanded Batch')\n",
        "]\n",
        "\n",
        "describe_distributions(poisson_distributions)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tfp.distributions.Poisson(\"One Poisson Scalar Batch/\", batch_shape=(), event_shape=(), dtype=float32)\n",
            "tfp.distributions.Poisson(\"Three Poissons/\", batch_shape=(3,), event_shape=(), dtype=float32)\n",
            "tfp.distributions.Poisson(\"Two-by-Three Poissons/\", batch_shape=(2, 3), event_shape=(), dtype=float32)\n",
            "tfp.distributions.Poisson(\"One Poisson Vector Batch/\", batch_shape=(1,), event_shape=(), dtype=float32)\n",
            "tfp.distributions.Poisson(\"One Poisson Expanded Batch/\", batch_shape=(1, 1), event_shape=(), dtype=float32)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lVPVIsC9rG4a",
        "colab_type": "text"
      },
      "source": [
        "The Poisson distribution is a scalar distribution, so its event shape is always `[]`. If we specify more rates, these show up in the batch shape. The final pair of examples is interesting: there's only a single rate, but because that rate is embedded in a numpy array with non-empty shape, that shape becomes the batch shape."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cFlXG9O5rG4b",
        "colab_type": "text"
      },
      "source": [
        "The standard Normal distribution is also a scalar. It's event shape is `[]`, just like for the Poisson, but we'll play with it to see our first example of *broadcasting*. The Normal is specified using `loc` and `scale` parameters:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e5PXRPM1rG4c",
        "colab_type": "code",
        "outputId": "928723a4-5547-455a-d121-e065e6cafbbe",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 85
        }
      },
      "source": [
        "normal_distributions = [\n",
        "    tfd.Normal(loc=0., scale=1., name='Standard'),\n",
        "    tfd.Normal(loc=[0.], scale=1., name='Standard Vector Batch'),\n",
        "    tfd.Normal(loc=[0., 1., 2., 3.], scale=1., name='Different Locs'),\n",
        "    tfd.Normal(loc=[0., 1., 2., 3.], scale=[[1.], [5.]],\n",
        "               name='Broadcasting Scale')\n",
        "]\n",
        "\n",
        "describe_distributions(normal_distributions)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tfp.distributions.Normal(\"Standard/\", batch_shape=(), event_shape=(), dtype=float32)\n",
            "tfp.distributions.Normal(\"Standard Vector Batch/\", batch_shape=(1,), event_shape=(), dtype=float32)\n",
            "tfp.distributions.Normal(\"Different Locs/\", batch_shape=(4,), event_shape=(), dtype=float32)\n",
            "tfp.distributions.Normal(\"Broadcasting Scale/\", batch_shape=(2, 4), event_shape=(), dtype=float32)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dh70eNXHrG4i",
        "colab_type": "text"
      },
      "source": [
        "The interesting example above is the `Broadcasting Scale` distribution. The `loc` parameter has shape `[4]`, and the `scale` parameter has shape `[2, 1]`. Using [Numpy broadcasting rules](https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html), the batch shape is `[2, 4]`. An equivalent (but less elegant and not-recommended) way to define the `\"Broadcasting Scale\"` distribution would be:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9G5JNBzQrG4j",
        "colab_type": "code",
        "outputId": "636b35de-3760-4598-be12-18f3f65decdb",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "describe_distributions(\n",
        "    [tfd.Normal(loc=[[0., 1., 2., 3], [0., 1., 2., 3.]],\n",
        "                scale=[[1., 1., 1., 1.], [5., 5., 5., 5.]])])"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tfp.distributions.Normal(\"Normal/\", batch_shape=(2, 4), event_shape=(), dtype=float32)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_hSBWsokrG4p",
        "colab_type": "text"
      },
      "source": [
        "We can see why the broadcasting notation is useful, although it's also a source of headaches and bugs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "trGxojHwrG4r",
        "colab_type": "text"
      },
      "source": [
        "### Sampling Scalar Distributions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TDJqRz-qrG4t",
        "colab_type": "text"
      },
      "source": [
        "There are two main things we can do with distributions: we can `sample` from them and we can compute `log_prob`s. Let's explore sampling first. The basic rule is that when we sample from a distribution, the resulting Tensor has shape `[sample_shape, batch_shape, event_shape]`, where `batch_shape` and `event_shape` are provided by the `Distribution` object, and `sample_shape` is provided by the call to `sample`. For scalar distributions, `event_shape = []`, so the Tensor returned from sample will have shape `[sample_shape, batch_shape]`. Let's try it:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2TbeP0btrG4u",
        "colab_type": "code",
        "outputId": "4e116872-b0e2-42fa-9fbe-c200ceb2c793",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 867
        }
      },
      "source": [
        "def describe_sample_tensor_shape(sample_shape, distribution):\n",
        "    print('Sample shape:', sample_shape)\n",
        "    print('Returned sample tensor shape:',\n",
        "          distribution.sample(sample_shape).shape)\n",
        "\n",
        "def describe_sample_tensor_shapes(distributions, sample_shapes):\n",
        "    started = False\n",
        "    for distribution in distributions:\n",
        "      print(distribution)\n",
        "      for sample_shape in sample_shapes:\n",
        "        describe_sample_tensor_shape(sample_shape, distribution)\n",
        "      print()\n",
        "\n",
        "sample_shapes = [1, 2, [1, 5], [3, 4, 5]]\n",
        "describe_sample_tensor_shapes(poisson_distributions, sample_shapes)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tfp.distributions.Poisson(\"One Poisson Scalar Batch/\", batch_shape=(), event_shape=(), dtype=float32)\n",
            "Sample shape: 1\n",
            "Returned sample tensor shape: (1,)\n",
            "Sample shape: 2\n",
            "Returned sample tensor shape: (2,)\n",
            "Sample shape: [1, 5]\n",
            "Returned sample tensor shape: (1, 5)\n",
            "Sample shape: [3, 4, 5]\n",
            "Returned sample tensor shape: (3, 4, 5)\n",
            "\n",
            "tfp.distributions.Poisson(\"Three Poissons/\", batch_shape=(3,), event_shape=(), dtype=float32)\n",
            "Sample shape: 1\n",
            "Returned sample tensor shape: (1, 3)\n",
            "Sample shape: 2\n",
            "Returned sample tensor shape: (2, 3)\n",
            "Sample shape: [1, 5]\n",
            "Returned sample tensor shape: (1, 5, 3)\n",
            "Sample shape: [3, 4, 5]\n",
            "Returned sample tensor shape: (3, 4, 5, 3)\n",
            "\n",
            "tfp.distributions.Poisson(\"Two-by-Three Poissons/\", batch_shape=(2, 3), event_shape=(), dtype=float32)\n",
            "Sample shape: 1\n",
            "Returned sample tensor shape: (1, 2, 3)\n",
            "Sample shape: 2\n",
            "Returned sample tensor shape: (2, 2, 3)\n",
            "Sample shape: [1, 5]\n",
            "Returned sample tensor shape: (1, 5, 2, 3)\n",
            "Sample shape: [3, 4, 5]\n",
            "Returned sample tensor shape: (3, 4, 5, 2, 3)\n",
            "\n",
            "tfp.distributions.Poisson(\"One Poisson Vector Batch/\", batch_shape=(1,), event_shape=(), dtype=float32)\n",
            "Sample shape: 1\n",
            "Returned sample tensor shape: (1, 1)\n",
            "Sample shape: 2\n",
            "Returned sample tensor shape: (2, 1)\n",
            "Sample shape: [1, 5]\n",
            "Returned sample tensor shape: (1, 5, 1)\n",
            "Sample shape: [3, 4, 5]\n",
            "Returned sample tensor shape: (3, 4, 5, 1)\n",
            "\n",
            "tfp.distributions.Poisson(\"One Poisson Expanded Batch/\", batch_shape=(1, 1), event_shape=(), dtype=float32)\n",
            "Sample shape: 1\n",
            "Returned sample tensor shape: (1, 1, 1)\n",
            "Sample shape: 2\n",
            "Returned sample tensor shape: (2, 1, 1)\n",
            "Sample shape: [1, 5]\n",
            "Returned sample tensor shape: (1, 5, 1, 1)\n",
            "Sample shape: [3, 4, 5]\n",
            "Returned sample tensor shape: (3, 4, 5, 1, 1)\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qiJK8UBorG40",
        "colab_type": "code",
        "outputId": "08fecdd6-7d68-4979-e75a-ff7dab498f0a",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 697
        }
      },
      "source": [
        "describe_sample_tensor_shapes(normal_distributions, sample_shapes)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tfp.distributions.Normal(\"Standard/\", batch_shape=(), event_shape=(), dtype=float32)\n",
            "Sample shape: 1\n",
            "Returned sample tensor shape: (1,)\n",
            "Sample shape: 2\n",
            "Returned sample tensor shape: (2,)\n",
            "Sample shape: [1, 5]\n",
            "Returned sample tensor shape: (1, 5)\n",
            "Sample shape: [3, 4, 5]\n",
            "Returned sample tensor shape: (3, 4, 5)\n",
            "\n",
            "tfp.distributions.Normal(\"Standard Vector Batch/\", batch_shape=(1,), event_shape=(), dtype=float32)\n",
            "Sample shape: 1\n",
            "Returned sample tensor shape: (1, 1)\n",
            "Sample shape: 2\n",
            "Returned sample tensor shape: (2, 1)\n",
            "Sample shape: [1, 5]\n",
            "Returned sample tensor shape: (1, 5, 1)\n",
            "Sample shape: [3, 4, 5]\n",
            "Returned sample tensor shape: (3, 4, 5, 1)\n",
            "\n",
            "tfp.distributions.Normal(\"Different Locs/\", batch_shape=(4,), event_shape=(), dtype=float32)\n",
            "Sample shape: 1\n",
            "Returned sample tensor shape: (1, 4)\n",
            "Sample shape: 2\n",
            "Returned sample tensor shape: (2, 4)\n",
            "Sample shape: [1, 5]\n",
            "Returned sample tensor shape: (1, 5, 4)\n",
            "Sample shape: [3, 4, 5]\n",
            "Returned sample tensor shape: (3, 4, 5, 4)\n",
            "\n",
            "tfp.distributions.Normal(\"Broadcasting Scale/\", batch_shape=(2, 4), event_shape=(), dtype=float32)\n",
            "Sample shape: 1\n",
            "Returned sample tensor shape: (1, 2, 4)\n",
            "Sample shape: 2\n",
            "Returned sample tensor shape: (2, 2, 4)\n",
            "Sample shape: [1, 5]\n",
            "Returned sample tensor shape: (1, 5, 2, 4)\n",
            "Sample shape: [3, 4, 5]\n",
            "Returned sample tensor shape: (3, 4, 5, 2, 4)\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wDRB80oLrG48",
        "colab_type": "text"
      },
      "source": [
        "That's about all there is to say about `sample`: returned sample tensors have shape `[sample_shape, batch_shape, event_shape]`.\n",
        "\n",
        "### Computing `log_prob` For Scalar Distributions\n",
        "\n",
        "Now let's take a look at `log_prob`, which is somewhat trickier. `log_prob` takes as input a (non-empty) tensor representing the location(s) at which to compute the `log_prob` for the distribution. In the most straightforward case, this tensor will have a shape of the form `[sample_shape, batch_shape, event_shape]`, where `batch_shape` and `event_shape` match the batch and event shapes of the distribution. Recall once more that for scalar distributions, `event_shape = []`, so the input tensor has shape `[sample_shape, batch_shape]` In this case, we get back a tensor of shape `[sample_shape, batch_shape]`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UgNIiFf9rG49",
        "colab_type": "code",
        "outputId": "c0b2f4e0-90b6-41af-b7c9-11aa9f0af483",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 54
        }
      },
      "source": [
        "three_poissons = tfd.Poisson(rate=[1., 10., 100.], name='Three Poissons')\n",
        "three_poissons"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tfp.distributions.Poisson 'Three Poissons/' batch_shape=(3,) event_shape=() dtype=float32>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OpN5WGog0WwC",
        "colab_type": "code",
        "outputId": "57d25cb5-b6fb-49ad-93a8-2b3bc32f62d6",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 88
        }
      },
      "source": [
        "three_poissons.log_prob([[1., 10., 100.], [100., 10., 1]])  # sample_shape is [2]."
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=594, shape=(2, 3), dtype=float32, numpy=\n",
              "array([[  -1.       ,   -2.0785608,   -3.222351 ],\n",
              "       [-364.73938  ,   -2.0785608,  -95.39483  ]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4szFj9lkrG5F",
        "colab_type": "code",
        "outputId": "dde8edfa-4899-4198-b144-62600d29a488",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 88
        }
      },
      "source": [
        "three_poissons.log_prob([[[[1., 10., 100.], [100., 10., 1.]]]])  # sample_shape is [1, 1, 2]."
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=601, shape=(1, 1, 2, 3), dtype=float32, numpy=\n",
              "array([[[[  -1.       ,   -2.0785608,   -3.222351 ],\n",
              "         [-364.73938  ,   -2.0785608,  -95.39483  ]]]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VG_n9BHsrG5M",
        "colab_type": "text"
      },
      "source": [
        "Note how in the first example, the input and output have shape `[2, 3]` and in the second example they have shape `[1, 1, 2, 3]`.\n",
        "\n",
        "That would be all there was to say, if it weren't for broadcasting. Here are the rules once we take broadcasting into account. We describe it in full generality and note simplifications for scalar distributions:\n",
        "1. Define `n = len(batch_shape) + len(event_shape)`. (For scalar distributions, `len(event_shape)=0`.)\n",
        "2. If the input tensor `t` has fewer than `n` dimensions, pad its shape by adding dimensions of size `1` on the left until it has exactly `n` dimensions. Call the resulting tensor `t'`.\n",
        "3. Broadcast the `n` rightmost dimensions of `t'` against the `[batch_shape, event_shape]` of the distribution you're computing a `log_prob` for. In more detail: for the dimensions where `t'` already matches the distribution, do nothing, and for the dimensions where `t'` has a singleton, replicate that singleton the appropriate number of times. Any other situation is an error. (For scalar distributions, we only broadcast against `batch_shape`, since event_shape = `[]`.)\n",
        "4. Now we're finally able to compute the `log_prob`. The resulting tensor will have shape `[sample_shape, batch_shape]`, where `sample_shape` is defined to be any dimensions of `t` or `t'` to the left of the `n`-rightmost dimensions: `sample_shape = shape(t)[:-n]`.\n",
        "\n",
        "This might be a mess if you don't know what it means, so let's work some examples:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YwDVaeRHrG5O",
        "colab_type": "code",
        "outputId": "181c8f91-07bd-4d0f-cbb4-ed4b7a068d75",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 54
        }
      },
      "source": [
        "three_poissons.log_prob([10.])"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=608, shape=(3,), dtype=float32, numpy=array([-16.104412 ,  -2.0785608, -69.052704 ], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xAImEhtdrG5U",
        "colab_type": "text"
      },
      "source": [
        "The tensor `[10.]` (with shape `[1]`) is broadcast across the `batch_shape` of 3, so we evaluate all three Poissons' log probability at the value 10."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "daDAG6p2rG5V",
        "colab_type": "code",
        "outputId": "dccd4e5f-8330-44c3-bedb-966f12c39b07",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 139
        }
      },
      "source": [
        "three_poissons.log_prob([[[1.], [10.]], [[100.], [1000.]]])"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=615, shape=(2, 2, 3), dtype=float32, numpy=\n",
              "array([[[-1.0000000e+00, -7.6974149e+00, -9.5394829e+01],\n",
              "        [-1.6104412e+01, -2.0785608e+00, -6.9052704e+01]],\n",
              "\n",
              "       [[-3.6473938e+02, -1.4348087e+02, -3.2223511e+00],\n",
              "        [-5.9131279e+03, -3.6195427e+03, -1.4069575e+03]]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 14
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "REEX-DgBrG5b",
        "colab_type": "text"
      },
      "source": [
        "In the above example, the input tensor has shape `[2, 2, 1]`, while the distributions object has a batch shape of 3. So for each of the `[2, 2]` sample dimensions, the single value provided gets broadcats to each of the three Poissons.\n",
        "\n",
        "A possibly useful way to think of it: because `three_poissons` has `batch_shape = [2, 3]`, a call to `log_prob` must take a Tensor whose last dimension is either 1 or 3; anything else is an error. (The numpy broadcasting rules treat the special case of a scalar as being totally equivalent to a Tensor of shape `[1]`.)\n",
        "\n",
        "Let's test our chops by playing with the more complex Poisson distribution with `batch_shape = [2, 3]`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MkSWkwYarG5d",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "poisson_2_by_3 = tfd.Poisson(\n",
        "    rate=[[1., 10., 100.,], [2., 20., 200.]],\n",
        "    name='Two-by-Three Poissons')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9YFRkkssrG5f",
        "colab_type": "code",
        "outputId": "dba707ce-63e4-406a-e8d0-6aadcfc7b341",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 88
        }
      },
      "source": [
        "poisson_2_by_3.log_prob(1.)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=623, shape=(2, 3), dtype=float32, numpy=\n",
              "array([[  -1.       ,   -7.697415 ,  -95.39483  ],\n",
              "       [  -1.3068528,  -17.004269 , -194.70168  ]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CqQXvOexrG5i",
        "colab_type": "code",
        "outputId": "d705598a-610f-41d5-afae-1f6f5a890875",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 88
        }
      },
      "source": [
        "poisson_2_by_3.log_prob([1.])  # Exactly equivalent to above, demonstrating the scalar special case."
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=630, shape=(2, 3), dtype=float32, numpy=\n",
              "array([[  -1.       ,   -7.697415 ,  -95.39483  ],\n",
              "       [  -1.3068528,  -17.004269 , -194.70168  ]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 17
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1nCuYQC5rG5m",
        "colab_type": "code",
        "outputId": "6cb0f831-1daf-45af-aaeb-8223e52b0ce2",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 88
        }
      },
      "source": [
        "poisson_2_by_3.log_prob([[1., 1., 1.], [1., 1., 1.]])  # Another way to write the same thing. No broadcasting."
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=637, shape=(2, 3), dtype=float32, numpy=\n",
              "array([[  -1.       ,   -7.697415 ,  -95.39483  ],\n",
              "       [  -1.3068528,  -17.004269 , -194.70168  ]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2PgG6udBrG5p",
        "colab_type": "code",
        "outputId": "1f577f2b-e7e8-4e2f-8627-4b498f514f9b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 88
        }
      },
      "source": [
        "poisson_2_by_3.log_prob([[1., 10., 100.]])  # Input is [1, 3] broadcast to [2, 3]."
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=644, shape=(2, 3), dtype=float32, numpy=\n",
              "array([[ -1.       ,  -2.0785608,  -3.222351 ],\n",
              "       [ -1.3068528,  -5.14709  , -33.907654 ]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Gm7ejyoArG5s",
        "colab_type": "code",
        "outputId": "c3dbe1ba-994d-4331-959d-73c89655ee4f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 88
        }
      },
      "source": [
        "poisson_2_by_3.log_prob([[1., 10., 100.], [1., 10., 100.]])  # Equivalent to above. No broadcasting."
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=651, shape=(2, 3), dtype=float32, numpy=\n",
              "array([[ -1.       ,  -2.0785608,  -3.222351 ],\n",
              "       [ -1.3068528,  -5.14709  , -33.907654 ]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mVMSGVvGrG5w",
        "colab_type": "code",
        "outputId": "2294b821-3c2c-448a-8e75-6471c3dbe26f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 88
        }
      },
      "source": [
        "poisson_2_by_3.log_prob([[1., 1., 1.], [2., 2., 2.]])  # No broadcasting."
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=658, shape=(2, 3), dtype=float32, numpy=\n",
              "array([[  -1.       ,   -7.697415 ,  -95.39483  ],\n",
              "       [  -1.3068528,  -14.701683 , -190.09651  ]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 21
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OVEpi5QErG5z",
        "colab_type": "code",
        "outputId": "a0518217-9162-47a7-ec74-46b7e5b72ecf",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 88
        }
      },
      "source": [
        "poisson_2_by_3.log_prob([[1.], [2.]])  # Equivalent to above. Input shape [2, 1] broadcast to [2, 3]."
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=665, shape=(2, 3), dtype=float32, numpy=\n",
              "array([[  -1.       ,   -7.697415 ,  -95.39483  ],\n",
              "       [  -1.3068528,  -14.701683 , -190.09651  ]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 22
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZW2tApDGrG53",
        "colab_type": "text"
      },
      "source": [
        "The above examples involved broadcasting over the batch, but the sample shape was empty. Suppose we have a collection of values, and we want to get the log probability of each value at each point in the batch. We could do it manually:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "03DvnmK2rG53",
        "colab_type": "code",
        "outputId": "e955762a-05e3-4dcb-a8f9-7244715565a4",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 139
        }
      },
      "source": [
        "poisson_2_by_3.log_prob([[[1., 1., 1.], [1., 1., 1.]], [[2., 2., 2.], [2., 2., 2.]]])  # Input shape [2, 2, 3]."
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=672, shape=(2, 2, 3), dtype=float32, numpy=\n",
              "array([[[  -1.       ,   -7.697415 ,  -95.39483  ],\n",
              "        [  -1.3068528,  -17.004269 , -194.70168  ]],\n",
              "\n",
              "       [[  -1.6931472,   -6.087977 ,  -91.4828   ],\n",
              "        [  -1.3068528,  -14.701683 , -190.09651  ]]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 23
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XkpJQ0dJrG56",
        "colab_type": "text"
      },
      "source": [
        "Or we could let broadcasting handle the last batch dimension:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KJ6OsodCrG57",
        "colab_type": "code",
        "outputId": "52ab9ca2-2ac1-45e1-9df7-c0fd7c9634ad",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 139
        }
      },
      "source": [
        "poisson_2_by_3.log_prob([[[1.], [1.]], [[2.], [2.]]])  # Input shape [2, 2, 1]."
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=679, shape=(2, 2, 3), dtype=float32, numpy=\n",
              "array([[[  -1.       ,   -7.697415 ,  -95.39483  ],\n",
              "        [  -1.3068528,  -17.004269 , -194.70168  ]],\n",
              "\n",
              "       [[  -1.6931472,   -6.087977 ,  -91.4828   ],\n",
              "        [  -1.3068528,  -14.701683 , -190.09651  ]]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 24
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eZFx8pThrG5-",
        "colab_type": "text"
      },
      "source": [
        "We can also (perhaps somewhat less naturally) let broadcasting handle just the first batch dimension:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UoGs7GBSrG5_",
        "colab_type": "code",
        "outputId": "a0a30bad-291b-4fb0-cd25-071108761800",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 139
        }
      },
      "source": [
        "poisson_2_by_3.log_prob([[[1., 1., 1.]], [[2., 2., 2.]]])  # Input shape [2, 1, 3]."
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=686, shape=(2, 2, 3), dtype=float32, numpy=\n",
              "array([[[  -1.       ,   -7.697415 ,  -95.39483  ],\n",
              "        [  -1.3068528,  -17.004269 , -194.70168  ]],\n",
              "\n",
              "       [[  -1.6931472,   -6.087977 ,  -91.4828   ],\n",
              "        [  -1.3068528,  -14.701683 , -190.09651  ]]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 25
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cOP4OhGDrG6C",
        "colab_type": "text"
      },
      "source": [
        "Or we could let broadcasting handle *both* batch dimensions:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tnG2f4tZrG6E",
        "colab_type": "code",
        "outputId": "0ff98b73-2977-4e5f-a9b4-276325474177",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 139
        }
      },
      "source": [
        "poisson_2_by_3.log_prob([[[1.]], [[2.]]])  # Input shape [2, 1, 1]."
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=693, shape=(2, 2, 3), dtype=float32, numpy=\n",
              "array([[[  -1.       ,   -7.697415 ,  -95.39483  ],\n",
              "        [  -1.3068528,  -17.004269 , -194.70168  ]],\n",
              "\n",
              "       [[  -1.6931472,   -6.087977 ,  -91.4828   ],\n",
              "        [  -1.3068528,  -14.701683 , -190.09651  ]]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 26
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I1s1drAwrG6K",
        "colab_type": "text"
      },
      "source": [
        "The above worked fine when we had only two values we wanted, but suppose we had a long list of values we wanted to evaluate at every batch point. For that, the following notation, which adds extra dimensions of size 1 to the right side of the shape, is extremely useful:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oUxbYZN_rG6K",
        "colab_type": "code",
        "outputId": "1e051274-af7f-481c-b5a6-2b1223650027",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 139
        }
      },
      "source": [
        "poisson_2_by_3.log_prob(tf.constant([1., 2.])[..., tf.newaxis, tf.newaxis])"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=704, shape=(2, 2, 3), dtype=float32, numpy=\n",
              "array([[[  -1.       ,   -7.697415 ,  -95.39483  ],\n",
              "        [  -1.3068528,  -17.004269 , -194.70168  ]],\n",
              "\n",
              "       [[  -1.6931472,   -6.087977 ,  -91.4828   ],\n",
              "        [  -1.3068528,  -14.701683 , -190.09651  ]]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 27
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Se893aIurG6M",
        "colab_type": "text"
      },
      "source": [
        "This is an instance of [strided slice notation](https://www.tensorflow.org/api_docs/cc/class/tensorflow/ops/strided-slice), which is worth knowing."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XNDhHqJmrG6N",
        "colab_type": "text"
      },
      "source": [
        "Going back to `three_poissons` for completeness, the same example looks like:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zKP7OmQsrG6N",
        "colab_type": "code",
        "outputId": "601907d5-3fc4-49de-9c6d-95b660144869",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 122
        }
      },
      "source": [
        "three_poissons.log_prob([[1.], [10.], [50.], [100.]])"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=711, shape=(4, 3), dtype=float32, numpy=\n",
              "array([[  -1.       ,   -7.697415 ,  -95.39483  ],\n",
              "       [ -16.104412 ,   -2.0785608,  -69.052704 ],\n",
              "       [-149.47777  ,  -43.34851  ,  -18.219254 ],\n",
              "       [-364.73938  , -143.48087  ,   -3.222351 ]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 28
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PK_9DwSdrG6R",
        "colab_type": "code",
        "outputId": "a69151a6-fcf4-4fb7-e8e6-731f5c85d881",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 122
        }
      },
      "source": [
        "three_poissons.log_prob(tf.constant([1., 10., 50., 100.])[..., tf.newaxis])  # Equivalent to above."
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=722, shape=(4, 3), dtype=float32, numpy=\n",
              "array([[  -1.       ,   -7.697415 ,  -95.39483  ],\n",
              "       [ -16.104412 ,   -2.0785608,  -69.052704 ],\n",
              "       [-149.47777  ,  -43.34851  ,  -18.219254 ],\n",
              "       [-364.73938  , -143.48087  ,   -3.222351 ]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 29
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lhL17DW5rG6T",
        "colab_type": "text"
      },
      "source": [
        "## Multivariate distributions\n",
        "\n",
        "We now turn to multivariate distributions, which have non-empty event shape. Let's look at multinomial distributions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lOdGa5n9rG6T",
        "colab_type": "code",
        "outputId": "ff169a97-36d5-4e12-b90c-c4d1f19736cd",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 85
        }
      },
      "source": [
        "multinomial_distributions = [\n",
        "    # Multinomial is a vector-valued distribution: if we have k classes,\n",
        "    # an individual sample from the distribution has k values in it, so the\n",
        "    # event_shape is `[k]`.\n",
        "    tfd.Multinomial(total_count=100., probs=[.5, .4, .1],\n",
        "                    name='One Multinomial'),\n",
        "    tfd.Multinomial(total_count=[100., 1000.], probs=[.5, .4, .1],\n",
        "                    name='Two Multinomials Same Probs'),\n",
        "    tfd.Multinomial(total_count=100., probs=[[.5, .4, .1], [.1, .2, .7]],\n",
        "                    name='Two Multinomials Same Counts'),\n",
        "    tfd.Multinomial(total_count=[100., 1000.],\n",
        "                    probs=[[.5, .4, .1], [.1, .2, .7]],\n",
        "                    name='Two Multinomials Different Everything')\n",
        "\n",
        "]\n",
        "\n",
        "describe_distributions(multinomial_distributions)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tfp.distributions.Multinomial(\"One Multinomial/\", batch_shape=(), event_shape=(3,), dtype=float32)\n",
            "tfp.distributions.Multinomial(\"Two Multinomials Same Probs/\", batch_shape=(2,), event_shape=(3,), dtype=float32)\n",
            "tfp.distributions.Multinomial(\"Two Multinomials Same Counts/\", batch_shape=(2,), event_shape=(3,), dtype=float32)\n",
            "tfp.distributions.Multinomial(\"Two Multinomials Different Everything/\", batch_shape=(2,), event_shape=(3,), dtype=float32)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-NQ8gK7irG6W",
        "colab_type": "text"
      },
      "source": [
        "Note how in the last three examples, the batch_shape is always `[2]`, but we can use broadcasting to either have a shared `total_count` or a shared `probs` (or neither), because under the hood they are broadcast to have the same shape.\n",
        "\n",
        "Sampling is straightforward, given what we know already:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hSr362qjrG6W",
        "colab_type": "code",
        "outputId": "94e87c96-8af5-423f-a0c6-6238d0b97fcc",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 697
        }
      },
      "source": [
        "describe_sample_tensor_shapes(multinomial_distributions, sample_shapes)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tfp.distributions.Multinomial(\"One Multinomial/\", batch_shape=(), event_shape=(3,), dtype=float32)\n",
            "Sample shape: 1\n",
            "Returned sample tensor shape: (1, 3)\n",
            "Sample shape: 2\n",
            "Returned sample tensor shape: (2, 3)\n",
            "Sample shape: [1, 5]\n",
            "Returned sample tensor shape: (1, 5, 3)\n",
            "Sample shape: [3, 4, 5]\n",
            "Returned sample tensor shape: (3, 4, 5, 3)\n",
            "\n",
            "tfp.distributions.Multinomial(\"Two Multinomials Same Probs/\", batch_shape=(2,), event_shape=(3,), dtype=float32)\n",
            "Sample shape: 1\n",
            "Returned sample tensor shape: (1, 2, 3)\n",
            "Sample shape: 2\n",
            "Returned sample tensor shape: (2, 2, 3)\n",
            "Sample shape: [1, 5]\n",
            "Returned sample tensor shape: (1, 5, 2, 3)\n",
            "Sample shape: [3, 4, 5]\n",
            "Returned sample tensor shape: (3, 4, 5, 2, 3)\n",
            "\n",
            "tfp.distributions.Multinomial(\"Two Multinomials Same Counts/\", batch_shape=(2,), event_shape=(3,), dtype=float32)\n",
            "Sample shape: 1\n",
            "Returned sample tensor shape: (1, 2, 3)\n",
            "Sample shape: 2\n",
            "Returned sample tensor shape: (2, 2, 3)\n",
            "Sample shape: [1, 5]\n",
            "Returned sample tensor shape: (1, 5, 2, 3)\n",
            "Sample shape: [3, 4, 5]\n",
            "Returned sample tensor shape: (3, 4, 5, 2, 3)\n",
            "\n",
            "tfp.distributions.Multinomial(\"Two Multinomials Different Everything/\", batch_shape=(2,), event_shape=(3,), dtype=float32)\n",
            "Sample shape: 1\n",
            "Returned sample tensor shape: (1, 2, 3)\n",
            "Sample shape: 2\n",
            "Returned sample tensor shape: (2, 2, 3)\n",
            "Sample shape: [1, 5]\n",
            "Returned sample tensor shape: (1, 5, 2, 3)\n",
            "Sample shape: [3, 4, 5]\n",
            "Returned sample tensor shape: (3, 4, 5, 2, 3)\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jjXgxPXCrG6Z",
        "colab_type": "text"
      },
      "source": [
        "Computing log probabilities is equally straightforward. Let's work an example with diagonal Multivariate Normal distributions. (Multinomials are not very broadcast friendly, since the constraints on the counts and probabilities mean broadcasting will often produce inadmissible values.) We'll use a batch of 2 3-dimensional distributions with the same mean but different scales (standard deviations):"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cnywBQdZrG6Z",
        "colab_type": "code",
        "outputId": "31a707eb-e272-4de8-bcc1-5297e91f2e7d",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "two_multivariate_normals = tfd.MultivariateNormalDiag(loc=[1., 2., 3.], scale_identity_multiplier=[1., 2.])\n",
        "two_multivariate_normals"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tfp.distributions.MultivariateNormalDiag 'MultivariateNormalDiag/' batch_shape=(2,) event_shape=(3,) dtype=float32>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 32
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S9xE21IirG6b",
        "colab_type": "text"
      },
      "source": [
        "(Note that although we used distributions where the scales were multiples of the identity, this is not a restriction on; we could pass `scale` instead of `scale_identity_multiplier`.)\n",
        "\n",
        "Now let's evaluate the log probability of each batch point at its mean and at a shifted mean:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YBOLH33PrG6b",
        "colab_type": "code",
        "outputId": "2a2cdc22-890a-471a-9c9c-c7f912ab2767",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        }
      },
      "source": [
        "two_multivariate_normals.log_prob([[[1., 2., 3.]], [[3., 4., 5.]]])  # Input has shape [2,1,3]."
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=2394, shape=(2, 2), dtype=float32, numpy=\n",
              "array([[-2.7568154, -4.836257 ],\n",
              "       [-8.756816 , -6.336257 ]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 33
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oPDC6y3qrG6e",
        "colab_type": "text"
      },
      "source": [
        "Exactly equivalently, we can use [https://www.tensorflow.org/api_docs/cc/class/tensorflow/ops/strided-slice](strided slice notation) to insert an extra shape=1 dimension in the middle of a constant:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "M9m9GMezrG6f",
        "colab_type": "code",
        "outputId": "3c531746-c695-44dc-96e8-a873cf052b43",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        }
      },
      "source": [
        "two_multivariate_normals.log_prob(\n",
        "    tf.constant([[1., 2., 3.], [3., 4., 5.]])[:, tf.newaxis, :])  # Equivalent to above."
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=2475, shape=(2, 2), dtype=float32, numpy=\n",
              "array([[-2.7568154, -4.836257 ],\n",
              "       [-8.756816 , -6.336257 ]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 34
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jJA07wN7rG6i",
        "colab_type": "text"
      },
      "source": [
        "On the other hand, if we don't insert the extra dimension, we pass `[1., 2., 3.]` to the first batch point and `[3., 4., 5.]` to the second:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xaX1unvPrG6i",
        "colab_type": "code",
        "outputId": "dc347782-0772-4eaa-e72b-5646ea0eeace",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "two_multivariate_normals.log_prob(tf.constant([[1., 2., 3.], [3., 4., 5.]]))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=2552, shape=(2,), dtype=float32, numpy=array([-2.7568154, -6.336257 ], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 35
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JnhW86vcUfT8",
        "colab_type": "text"
      },
      "source": [
        "## Shape Manipulation Techniques"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6EYcFW7OrG6m",
        "colab_type": "text"
      },
      "source": [
        "### The Reshape Bijector\n",
        "\n",
        "The `Reshape` bijector can be used to reshape the *event_shape* of a distribution. Let's see an example:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1YT_lQCarG6m",
        "colab_type": "code",
        "outputId": "316040e9-2470-4cf2-a711-b3a673124ed8",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "six_way_multinomial = tfd.Multinomial(total_count=1000., probs=[.3, .25, .2, .15, .08, .02])\n",
        "six_way_multinomial"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tfp.distributions.Multinomial 'Multinomial/' batch_shape=() event_shape=(6,) dtype=float32>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 36
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c5a5uXQsUpMs",
        "colab_type": "text"
      },
      "source": [
        "We created a multinomial with an event shape of `[6]`. The Reshape Bijector allows us to treat this as a distribution with an event shape of `[2, 3]`.\n",
        "\n",
        "A `Bijector` represents a differentiable, one-to-one function on an open subset of ${\\mathbb R}^n$. `Bijectors` are used in conjunction with `TransformedDistribution`, which models a distribution $p(y)$ in terms of a base distribution $p(x)$ and a `Bijector` that represents $Y = g(X)$.\n",
        "Let's see it in action:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Wttfn9Q-rG6o",
        "colab_type": "code",
        "outputId": "56d72b96-9c49-4117-c56a-3be1275ddf8c",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "transformed_multinomial = tfd.TransformedDistribution(\n",
        "    distribution=six_way_multinomial,\n",
        "    bijector=tfb.Reshape(event_shape_out=[2, 3]))\n",
        "transformed_multinomial"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tfp.distributions.TransformedDistribution 'reshapeMultinomial/' batch_shape=() event_shape=(2, 3) dtype=float32>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 37
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sh6l4XZdrG6p",
        "colab_type": "code",
        "outputId": "4b9bf391-3c83-469b-8e11-75ea29336eff",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "six_way_multinomial.log_prob([500., 100., 100., 150., 100., 50.])"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=2628, shape=(), dtype=float32, numpy=-178.21973>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 38
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6nMHlVrArG6r",
        "colab_type": "code",
        "outputId": "3f6b49ed-7e42-4c68-fb78-b2c47d2b6244",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "transformed_multinomial.log_prob([[500., 100., 100.], [150., 100., 50.]])"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=2703, shape=(), dtype=float32, numpy=-178.21973>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 39
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dxZNZ02OrG6t",
        "colab_type": "text"
      },
      "source": [
        "This is the *only* thing the `Reshape` bijector can do: it cannot turn event dimensions into batch dimensions or vice-versa."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "de7ek-FerG6t",
        "colab_type": "text"
      },
      "source": [
        "### The Independent Distribution\n",
        "\n",
        "The `Independent` distribution is used to treat a collection of independent, not-necessarily-identical (aka a batch of) distributions as a single distribution. More concisely, `Independent` allows to convert dimensions in `batch_shape` to dimensions in `event_shape`. We'll illustrate by example:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tLwioZPRrG6t",
        "colab_type": "code",
        "outputId": "38290806-baa5-472f-c75b-0eed509822ea",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "two_by_five_bernoulli = tfd.Bernoulli(\n",
        "    probs=[[.05, .1, .15, .2, .25], [.3, .35, .4, .45, .5]],\n",
        "    name=\"Two By Five Bernoulli\")\n",
        "two_by_five_bernoulli"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tfp.distributions.Bernoulli 'Two By Five Bernoulli/' batch_shape=(2, 5) event_shape=() dtype=int32>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 40
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-okVviR3rG6v",
        "colab_type": "text"
      },
      "source": [
        "We can think of this as two-by-five array of coins with the associated probabilities of heads. Let's evaluate the probability of a particular, arbitrary set of ones-and-zeros:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9yq9jTGIrG6x",
        "colab_type": "code",
        "outputId": "1de9e93f-0a76-45a8-e54a-99f0679de68e",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 85
        }
      },
      "source": [
        "pattern = [[1., 0., 0., 1., 0.], [0., 0., 1., 1., 1.]]\n",
        "two_by_five_bernoulli.log_prob(pattern)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=2722, shape=(2, 5), dtype=float32, numpy=\n",
              "array([[-2.9957323 , -0.10536051, -0.16251893, -1.609438  , -0.28768206],\n",
              "       [-0.35667494, -0.43078294, -0.9162907 , -0.7985077 , -0.6931472 ]],\n",
              "      dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 41
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C9CA19oPrG6y",
        "colab_type": "text"
      },
      "source": [
        "We can use `Independent` to turn this into two different \"sets of five Bernoulli's\", which is useful if we want to consider a \"row\" of coin flips coming up in a given pattern as a single outcome:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1iR23BMBrG6z",
        "colab_type": "code",
        "outputId": "9f8e80ef-23ce-48cf-b77c-eb3070ec255d",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "two_sets_of_five = tfd.Independent(\n",
        "    distribution=two_by_five_bernoulli,\n",
        "    reinterpreted_batch_ndims=1,\n",
        "    name=\"Two Sets Of Five\")\n",
        "two_sets_of_five"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tfp.distributions.Independent 'Two Sets Of Five/' batch_shape=(2,) event_shape=(5,) dtype=int32>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 42
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mRrkesaPrG67",
        "colab_type": "text"
      },
      "source": [
        "Mathematically, we're computing the log probability of each \"set\" of five by summing the log probabilities of the five \"independent\" coin flips in the set, which is where the distribution gets its name:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LcM6OgKNrG66",
        "colab_type": "code",
        "outputId": "afc999e6-fc41-438f-aa93-de1bcd09a974",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "two_sets_of_five.log_prob(pattern)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=2739, shape=(2,), dtype=float32, numpy=array([-5.160732 , -3.1954036], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 43
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "krpnUVL9rG7A",
        "colab_type": "text"
      },
      "source": [
        "We can go even further and use `Independent` to create a distribution where individual events are a set of two-by-five Bernoulli's:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PXSsoidirG7A",
        "colab_type": "code",
        "outputId": "01de8e83-ba2d-45a4-b540-672dabcc3e81",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "one_set_of_two_by_five = tfd.Independent(\n",
        "    distribution=two_by_five_bernoulli, reinterpreted_batch_ndims=2,\n",
        "    name=\"One Set Of Two By Five\")\n",
        "one_set_of_two_by_five.log_prob(pattern)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: id=2756, shape=(), dtype=float32, numpy=-8.356134>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 44
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QfbfHA4hrG7F",
        "colab_type": "text"
      },
      "source": [
        "It's worth noting that from the perspective of `sample`, using `Independent` changes nothing:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uZ3NhQEZrG7F",
        "colab_type": "code",
        "outputId": "94426140-6c0b-44de-cd21-20dc0593f86d",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 221
        }
      },
      "source": [
        "describe_sample_tensor_shapes(\n",
        "    [two_by_five_bernoulli,\n",
        "     two_sets_of_five,\n",
        "     one_set_of_two_by_five],\n",
        "    [[3, 5]])"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "tfp.distributions.Bernoulli(\"Two By Five Bernoulli/\", batch_shape=(2, 5), event_shape=(), dtype=int32)\n",
            "Sample shape: [3, 5]\n",
            "Returned sample tensor shape: (3, 5, 2, 5)\n",
            "\n",
            "tfp.distributions.Independent(\"Two Sets Of Five/\", batch_shape=(2,), event_shape=(5,), dtype=int32)\n",
            "Sample shape: [3, 5]\n",
            "Returned sample tensor shape: (3, 5, 2, 5)\n",
            "\n",
            "tfp.distributions.Independent(\"One Set Of Two By Five/\", batch_shape=(), event_shape=(2, 5), dtype=int32)\n",
            "Sample shape: [3, 5]\n",
            "Returned sample tensor shape: (3, 5, 2, 5)\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "usTT7v0trG7H",
        "colab_type": "text"
      },
      "source": [
        "As a parting exercise for the reader, we suggest considering the differences and similarities between a vector batch of `Normal` distributions and a `MultivariateNormalDiag` distribution from a sampling and log probability perspective. How can we use `Independent` to construct a `MultivariateNormalDiag` from a batch of `Normal`s? (Note that `MultivariateNormalDiag` is not actually implemented this way.)"
      ]
    }
  ]
}